"""
浏览器MCP客户端
基于Playwright MCP服务器提供浏览器自动化功能
"""

import asyncio
from typing import Dict, Any, List

from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
from langchain_mcp_adapters.tools import load_mcp_tools
from mcp import StdioServerParameters, stdio_client, ClientSession
from langgraph.prebuilt import create_react_agent

from app.bailian.common import llm


class BrowserMCPClient:
    """浏览器MCP客户端类"""
    
    def __init__(self):
        self.session = None
        self.tools = None
        self.agent = None
    
    async def initialize(self):
        """初始化浏览器MCP客户端"""
        # 使用npx启动远程MCP服务器
        server_params = StdioServerParameters(
            command="npx",
            args=["-y", "@executeautomation/playwright-mcp-server"]
        )
        
        async with stdio_client(server_params) as (read, write):
            async with ClientSession(read, write) as session:
                self.session = session
                await session.initialize()
                # 获取MCP工具
                self.tools = await load_mcp_tools(session)
                print("可用的浏览器工具:", [tool.name for tool in self.tools])
                
                # 创建React代理
                self.agent = create_react_agent(model=llm, tools=self.tools, debug=True)
                return self
    
    async def browse_website(self, url: str, action: str = "获取页面内容") -> Dict[str, Any]:
        """浏览网站并执行操作"""
        if not self.agent:
            await self.initialize()
        
        prompt = f"请访问 {url} 并{action}"
        response = await self.agent.ainvoke(input={"messages": [("user", prompt)]})
        
        return self._parse_response(response)
    
    async def search_and_extract(self, search_query: str, website: str = "百度", extract_info: str = "相关信息") -> Dict[str, Any]:
        """搜索并提取信息"""
        if not self.agent:
            await self.initialize()
        
        prompt = f"请在{website}中搜索'{search_query}'，并提取{extract_info}"
        response = await self.agent.ainvoke(input={"messages": [("user", prompt)]})
        
        return self._parse_response(response)
    
    async def fill_form_and_submit(self, url: str, form_data: Dict[str, str], submit_button: str = "提交") -> Dict[str, Any]:
        """填写表单并提交"""
        if not self.agent:
            await self.initialize()
        
        form_info = ", ".join([f"{key}: {value}" for key, value in form_data.items()])
        prompt = f"请访问 {url}，填写表单({form_info})，并点击{submit_button}按钮"
        response = await self.agent.ainvoke(input={"messages": [("user", prompt)]})
        
        return self._parse_response(response)
    
    def _parse_response(self, response: Dict[str, Any]) -> Dict[str, Any]:
        """解析代理响应"""
        result = {
            "messages": [],
            "tool_calls": [],
            "final_answer": ""
        }
        
        messages = response["messages"]
        for message in messages:
            if isinstance(message, HumanMessage):
                result["messages"].append({"role": "user", "content": message.content})
            elif isinstance(message, AIMessage):
                if message.content:
                    result["final_answer"] = message.content
                    result["messages"].append({"role": "assistant", "content": message.content})
                else:
                    for tool_call in message.tool_calls:
                        result["tool_calls"].append({
                            "name": tool_call["name"],
                            "args": tool_call["args"]
                        })
            elif isinstance(message, ToolMessage):
                result["messages"].append({"role": "tool", "name": message.name, "content": message.content})
        
        return result


async def demo_browser_mcp():
    """演示浏览器MCP功能"""
    browser_client = BrowserMCPClient()
    await browser_client.initialize()
    
    # 演示1: 浏览网站
    print("=== 演示1: 浏览百度首页 ===")
    result1 = await browser_client.browse_website("https://www.baidu.com", "获取页面标题和主要内容")
    print("结果:", result1["final_answer"])
    
    # 演示2: 搜索信息
    print("\n=== 演示2: 搜索天气信息 ===")
    result2 = await browser_client.search_and_extract("北京今天天气", "百度", "温度和天气情况")
    print("结果:", result2["final_answer"])
    
    # 演示3: 表单操作（示例）
    print("\n=== 演示3: 表单操作 ===")
    # 注意：实际使用时需要替换为真实的表单页面
    # result3 = await browser_client.fill_form_and_submit(
    #     "https://example.com/login",
    #     {"username": "test", "password": "test123"}
    # )
    # print("结果:", result3["final_answer"])


if __name__ == "__main__":
    asyncio.run(demo_browser_mcp())